package com.ly.robustvideomatting.model;

import org.opencv.core.*;
import org.opencv.imgcodecs.Imgcodecs;
import org.opencv.imgproc.Imgproc;

import java.util.ArrayList;
import java.util.List;

public class ImageMatting {
    public static Mat applyMask(Mat originalImage, Mat mask) {
        return matting(originalImage, mask);
    }

    private static Mat matting(Mat originalImage, Mat mask) {
        // 确保原始图像和蒙版尺寸相同
        if (!originalImage.size().equals(mask.size())) {
            throw new IllegalArgumentException("Original image and mask must have the same size");
        }

        // 如果蒙版不是单通道，将其转换为单通道
        Mat singleChannelMask = new Mat();
        if (mask.channels() != 1) {
            Imgproc.cvtColor(mask, singleChannelMask, Imgproc.COLOR_BGR2GRAY);
        } else {
            mask.copyTo(singleChannelMask);
        }

        // 创建输出图像
        Mat result = new Mat(originalImage.size(), CvType.CV_8UC4);

        // 将原始图像转换为4通道（如果不是的话）
        Mat originalImageRGBA = new Mat();
        if (originalImage.channels() == 3) {
            Imgproc.cvtColor(originalImage, originalImageRGBA, Imgproc.COLOR_BGR2BGRA);
        } else if (originalImage.channels() == 4) {
            originalImage.copyTo(originalImageRGBA);
        } else {
            throw new IllegalArgumentException("Original image must have 3 or 4 channels");
        }

        // 归一化蒙版到0-1范围
        Mat normalizedMask = new Mat();
        singleChannelMask.convertTo(normalizedMask, CvType.CV_32F, 1.0/255.0);

        // 分离原始图像的通道
        List<Mat> channels = new ArrayList<>();
        Core.split(originalImageRGBA, channels);

        // 应用蒙版到每个通道
        for (int i = 0; i < 3; i++) {  // 只处理BGR通道，保持Alpha通道不变
            Mat temp = new Mat();
            Core.multiply(channels.get(i), normalizedMask, temp, 1, CvType.CV_8U);
            channels.set(i, temp);
        }

        // 设置Alpha通道
        Mat alphaChannel = new Mat();
        normalizedMask.convertTo(alphaChannel, CvType.CV_8U, 255);
        channels.set(3, alphaChannel);

        // 合并通道
        Core.merge(channels, result);

        return result;
    }

    public static void main(String[] args) {
//        nu.pattern.OpenCV.loadLocally();
        String imagePath = "output.png";  // 原始图像路径
        String maskPath = "output1.png";    // 蒙版图像路径

        Mat imread = Imgcodecs.imread(imagePath);
        Mat mask = Imgcodecs.imread(maskPath);
        long l = System.currentTimeMillis();
        Mat mat = applyMask(imread, mask);
        System.out.println(System.currentTimeMillis() - l);
        Imgcodecs.imwrite("out.png",mat);
    }
}